from PIL import Image
import win32gui
import pyautogui
import mss
import cv2
import numpy as np
from core.log import logger
import setting as c

hwnd_title = dict()

def get_all_hwnd(hwnd,mouse):
    if win32gui.IsWindow(hwnd) and win32gui.IsWindowEnabled(hwnd) and win32gui.IsWindowVisible(hwnd):
        hwnd_title.update({hwnd:win32gui.GetWindowText(hwnd)})

def get_screenshot_ndarray(hwnd) ->np.ndarray:
    rect = win32gui.GetWindowRect(hwnd)  
    img = pyautogui.screenshot() 
    img = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)
    img = img[rect[1]:rect[3], rect[0]:rect[2]]
    
    return img

def get_screenshot_image(hwnd) -> Image:
    rect = win32gui.GetWindowRect(hwnd)  
    img = pyautogui.screenshot() 
    rect = (rect[0], rect[1], rect[2], rect[3])
    img = img.crop(rect)
    return img

def mss_screenshot(monitor_number, rect):
    with mss.mss() as sct:
        mon = sct.monitors[monitor_number]
        monitor = {
            "top": rect[1],
            "left": rect[0],
            "width": rect[2] - rect[0],
            "height": rect[3] - rect[1],
            "mon": monitor_number,
        }
        output = "sct-mon{mon}_{top}x{left}_{width}x{height}.png".format(**monitor)

        # Grab the data
        sct_img = sct.grab(monitor)
        # Save to the picture file
        # mss.tools.to_png(sct_img.rgb, sct_img.size, output=output)
        # print(output)
        
        # Convert to PIL image for more complex manipulation (optional)
        img = Image.frombytes('RGB', sct_img.size, sct_img.rgb)
    
        return img
    
def get_mss_screenshot_image(hwnd) -> Image:
    rect = win32gui.GetWindowRect(hwnd)  
    logger.info(f"mss rect: {rect} screen:{c.screen_width} x {c.screen_height}")
    monitor_number = 1
    if rect[0] > c.screen_width:
        monitor_number = 2
    img = mss_screenshot(monitor_number, rect) 
    return img